/***************************************************************************
 *
 *   Copyright (C) 2010 - 2024 by Willem van Straten
 *   Licensed under the Academic Free License version 2.1
 *
 ***************************************************************************/

#include "dsp/TransferCUDA.h"
#include "dsp/WeightedTimeSeries.h"
#include "dsp/MemoryCUDA.h"

#include "Error.h"

#include <iostream>
#include <string.h>
#include <assert.h>

using namespace std;

void xfer (void* dst, const void* src, size_t count, cudaMemcpyKind kind, cudaStream_t stream = 0)
{
  assert (dst != nullptr);
  assert (src != nullptr);
  assert (count > 0);

  cudaError error;
  if (stream)
    error = cudaMemcpyAsync (dst, src, count, kind, stream);
  else
    error = cudaMemcpy (dst, src, count, kind);

  if (error == cudaSuccess)
    return;

  if (dsp::Operation::verbose)
    std::cerr << "dsp::TransferCUDA::transformation failed with error " << cudaGetErrorString(error) << std::endl;

  throw Error (InvalidState, "dsp::TransferCUDA::transformation", cudaGetErrorString (error));
}

//! Default constructor- always inplace
dsp::TransferCUDA::TransferCUDA(cudaStream_t _stream)
  : Transformation<TimeSeries,TimeSeries> ("CUDA::Transfer", outofplace)
{
  stream = _stream;
  kind = cudaMemcpyHostToDevice;
}

//! Do stuff
void dsp::TransferCUDA::transformation ()
{
  if (verbose)
    cerr << "dsp::TransferCUDA::transformation" << endl;

  prepare ();

  // only transfer data if there is valid data to transfer
  if (input->get_ndat() == 0)
  {
    if (verbose)
      cerr << "dsp::TransferCUDA::transformation skipping transfer as ndat=" << input->get_ndat() << endl;
    return;
  }

  if (kind == cudaMemcpyHostToDevice)
  {
    if (verbose)
      cerr << "dsp::TransferCUDA::transformation synchronizing pre H2D transfer" << endl;
    if (stream)
      cudaStreamSynchronize(stream);
    else
      cudaDeviceSynchronize();
  }

  if (verbose)
  {
    cerr << "dsp::TransferCUDA::transformation input ndat="
         << input->get_ndat() << " ndim=" << input->get_ndim()
         << " nchan=" << input->get_nchan() << " npol=" << input->get_npol();
    if (input->get_order() == TimeSeries::OrderFPT)
    {
      if (input->get_npol() > 1)
        cerr << " span=" << input->get_datptr (0,1) - input->get_datptr(0,0);
      cerr << " offset=" << input->get_datptr(0,0) - (float*)input->internal_get_buffer() << endl;
    }
    else
      cerr << endl;
  }

  if (verbose)
    cerr << "dsp::TransferCUDA::transformation xfer data" << endl;

  xfer (output->internal_get_buffer(), input->internal_get_buffer(), input->internal_get_size(), kind, stream);

  auto weighted_input = dynamic_cast<const WeightedTimeSeries*>(input.get());
  auto weighted_output = dynamic_cast<WeightedTimeSeries*>(output.get());

  if (weighted_input && weighted_output)
  {
    auto input_device_memory = dynamic_cast<const CUDA::DeviceMemory*>(weighted_input->get_weights_memory());
    auto output_device_memory = dynamic_cast<CUDA::DeviceMemory*>(weighted_output->get_weights_memory());

    void* dst = weighted_output->internal_get_weights_buffer();
    const void* src = weighted_input->internal_get_weights_buffer();
    size_t count = weighted_input->internal_get_weights_size();

    assert (weighted_output->internal_get_weights_size() >= count);

    if (kind == cudaMemcpyHostToDevice)
    {
      if (input_device_memory != nullptr)
        throw (InvalidState, "dsp::TransferCUDA::transformation",
              "WeightedTimeSeries input weights are not in host memory (host->device)");

      if (output_device_memory != nullptr)
      {
        // input weights are on host and output weights are on device
        if (verbose)
          cerr << "dsp::TransferCUDA::transformation xfer weights host->device" << endl;
        xfer(dst, src, count, kind, stream);
      }
      else
      {
        // both input and output weights are on host
        memcpy(dst, src, count);
      }
    }
    else if (kind == cudaMemcpyDeviceToHost)
    {
      if (output_device_memory != nullptr)
        throw (InvalidState, "dsp::TransferCUDA::transformation",
              "WeightedTimeSeries output weights are not in host memory (device->host)");

      if (input_device_memory != nullptr)
      {
        // input weights are on device and output weights are on host
        if (verbose)
          cerr << "dsp::TransferCUDA::transformation xfer weights device->host" << endl;
        xfer(dst, src, count, kind, stream);
      }
      else
      {
        // both input and output weights are on host
        memcpy(dst, src, count);
      }
    }
  }

  if (verbose)
  {
    cerr << "dsp::TransferCUDA::transformation output ndat="
         << output->get_ndat() << " ndim=" << output->get_ndim()
         << " nchan=" << output->get_nchan() << " npol=" << output->get_npol();
    if (output->get_order() == TimeSeries::OrderFPT)
    {
      if (output->get_npol() > 1)
        cerr << " span=" << output->get_datptr (0, 1) - output->get_datptr(0,0);
      cerr << " offset=" << output->get_datptr(0,0) - (float*)output->internal_get_buffer() << endl;
    }
    else
      cerr << endl;
  }

  if (kind == cudaMemcpyDeviceToHost)
  {
    if (verbose)
      cerr << "dsp::TransferCUDA::transformation synchronizing after device-to-host transfer" << endl;
    if (stream)
      cudaStreamSynchronize(stream);
    else
      cudaDeviceSynchronize();
  }
}

void dsp::TransferCUDA::prepare ()
{
  output->set_match( const_cast<TimeSeries*>(input.get()) );
  output->internal_match( input );
  output->copy_configuration( input );

  if (verbose)
    std::cerr << "dsp::TransferCUDA::prepare input->ndat=" << input->get_ndat() << " output->ndat=" << output->get_ndat() << std::endl;
}
